import torch
from torch import nn
from transformers import GPT2Config, GPT2Model

# 定义一个与GPT-2相似的配置
class MyTransformerConfig(GPT2Config):
    model_type = "my_transformer"
    vocab_size = 50257  # 以GPT-2的vocab_size为例
    n_ctx = 1024        # 最大序列长度
    n_positions = 1024  # 同上
    n_embd = 768        # 词嵌入的维度
    n_layer = 12         # Transformer层数
    n_head = 12         # 注意力头数
    activation_function = "gelu"  # 激活函数
    resid_pdrop = 0.1   # 残差连接的dropout率
    embd_pdrop = 0.1    # 词嵌入的dropout率
    attn_pdrop = 0.1     # 注意力机制的dropout率
    layer_norm_epsilon = 1e-5  # LayerNorm的epsilon


class MyTransformerDecoder(GPT2Model):
    def __init__(self, config):
        super().__init__(config)
        self.transformer = GPT2Model(config).transformer  # 使用GPT2的Transformer层

    def forward(self, input_ids, past_key_values=None):
        # 只使用Transformer的解码器部分
        decoder_input_ids = input_ids[:, :-1].contiguous()  # 预测时移位的输入
        decoder_attention_mask = self.create_attention_mask(decoder_input_ids)

        # 通过Transformer层
        decoder_output = self.transformer.decoder(
            decoder_input_ids,
            past_key_values=past_key_values,
            attention_mask=decoder_attention_mask
        )

        # 只返回解码器的输出
        return decoder_output

class MyLMHead(nn.Module):
    def __init__(self, model_embeddings_weights, config):
        super().__init__()
        self.predictions = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.predictions.weight = model_embeddings_weights

    def forward(self, hidden_states):
        return self.predictions(hidden_states)


class MyGPT2(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = MyTransformerDecoder(config)
        self.lm_head = MyLMHead(self.transformer.model.model.embeddings.word_embeddings.weight, config)

    def forward(self, input_ids, past_key_values=None):
        # 调用解码器
        decoder_output = self.transformer.forward(input_ids, past_key_values=past_key_values)

        # 应用语言模型头
        logits = self.lm_head(decoder_output[0])

        return logits

config = MyTransformerConfig()
model = MyGPT2(config)

# 指定GPT-2模型的名称，你可以选择 'gpt2', 'gpt2-medium', 'gpt2-large' 等
model_name = "gpt2"

# 加载预训练的GPT-2模型
gpt2_model = GPT2Model.from_pretrained(model_name)

# 获取词嵌入权重，这些权重是模型状态字典中的一个键值对
model_embeddings = gpt2_model.transformer.wte.weight  # wte代表word token embeddings
# 准备输入，例如使用GPT-2的预训练词嵌入
# model_embeddings = ...  # 这里应该是GPT-2的词嵌入权重
input_ids = torch.tensor([1,2,3,4])  # 你的输入序列

# 将词嵌入权重传递给模型
model.lm_head.predictions.weight = model_embeddings

# 前向传播
logits = model(input_ids)


if __name__ == '__main__':
    print("==========over==========")